from keras.datasets import cifar10
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Activation, Flatten
from keras.layers.convolutional import Convolution2D, MaxPooling2D
from keras.optimizers import SGD,adam
from keras.utils import np_utils
from DataUtils import load_train32,load_val32_Set14
import numpy as np
from DataUtils import crop_center_board
from model import SRCNN, SRCNN_test
import os
from skimage import io,util,exposure,measure
import skimage
import matplotlib.pyplot as plt
import sys
import keras.backend as K
from Data_utils.Metrics import PSNR
import json
import tensorflow as tf
import numpy
import cv2

def custom_loss(y_true,y_pred):
    MSE = K.mean(K.square(y_pred - y_true))
    constraint = K.mean(K.abs(y_pred))
    loss = 5*MSE + constraint
    return loss

def PSNR_loss(I,P):
    """
    #sess = K.get_session()
    print 'I: {0}'.format(I)
    print 'I shape: {0}'.format(tf.shape(I))
    print 'I size: {0}'.format(tf.size(I))
    """
    batch_size = 128.
    I = K.reshape(I,(int(batch_size),18*18))
    P = K.reshape(P,(int(batch_size),18*18))
    MSE = K.mean(K.square(I-P),axis=1)
    PSNR_ = 20. * (tf.log(255./K.sqrt(MSE))/tf.log(10.))
    PSNR_ = K.mean(PSNR_)
    return -PSNR_

def train_model():
    train_X, train_Y = load_train32(factor=2.,input_size=32,stride=14)
    val_X, val_Y = load_val32_Set14(factor=2.,input_size=32,stride=14)
    train_Y = crop_center_board(train_Y,board_size=6)
    val_Y = crop_center_board(val_Y,board_size=6)
    print '{0} train samples'.format(len(train_Y))
    print '{0} val samples'.format(len(val_Y))

    train_X /= 255.
    train_Y /= 255.
    val_X /= 255.
    val_Y /= 255.

    with K.tf.device('/gpu:1'):
        K.set_session(K.tf.Session(config=K.tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)))
        model = SRCNN(img_channels=1,img_rows=32,img_cols=32)
        model.compile(optimizer=SGD(lr=0.01,momentum=0.9,nesterov=True),loss= 'mse')
    """
    hist = model.fit({'conv_input': train_X, 'skip_input': train_X, 'output': train_Y},
                     nb_epoch=10,
                     validation_data={'conv_input': val_X, 'skip_input': val_X, 'output': val_Y},
                     batch_size=128,shuffle=True)
    """
    hist = model.fit(X=train_X,y=train_Y,nb_epoch=10,validation_data=(val_X,val_Y),batch_size=128,shuffle=True)
    model.save_weights('metadata/sr_weights_mse_srcnn.h5',overwrite=True)
    training = {
        'loss': hist.history['loss'],
        'val_loss': hist.history['val_loss'],
        'optimizer': model.optimizer.get_config(),
    }
    f = open('metadata/sr_training_mse_srcnn.json', 'wb')
    meta_json = json.dumps(training, default=lambda o: o.__dict__, indent=4)
    f.write(meta_json)
    f.close()

def inference():
    factor = 2.
    folder = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/SRCNNset/Test/Set5'
    for img in os.listdir(folder):
        img_path = folder + '/' + img
        img_arr_RGB = io.imread(img_path)

        hei = img_arr_RGB.shape[1]
        wid = img_arr_RGB.shape[0]
        model = SRCNN_test(img_channels=1,img_cols=hei,img_rows=wid)
        model.load_weights('metadata/sr_weights_mse_srcnn.h5')
        #blurring
        img_arr_YUV = cv2.cvtColor(img_arr_RGB,cv2.COLOR_RGB2YUV)
        img_arr_UV = img_arr_YUV[:,:,1:]
        img_arr_Y_blur_ = cv2.resize(img_arr_YUV[:,:,0],None,fx=1./factor, fy=1./factor, interpolation = cv2.INTER_CUBIC)
        img_arr_Y_blur = cv2.resize(img_arr_Y_blur_,None,fx=factor, fy=factor, interpolation = cv2.INTER_CUBIC)
        img_arr_Y_blur = np.reshape(img_arr_Y_blur,(1,1,wid,hei))
        #generate recostucted YUV
        img_arr_SR_Y_ = K.eval(model(K.variable(img_arr_Y_blur/255.))) * 255.
        img_arr_SR_Y = np.reshape(img_arr_SR_Y_,(1,wid,hei)).transpose(1,2,0)
        img_arr_SR_YUV = np.dstack((img_arr_SR_Y,img_arr_UV))

        img_arr_blur_YUV = np.dstack((np.reshape(img_arr_Y_blur,(1,wid,hei)).transpose(1,2,0),img_arr_UV))
        #img_arr_blur_YUV = np.dstack((img_arr_UV,np.reshape(img_arr_Y_blur,(1,wid,hei)).transpose(1,2,0)))

        # YUV to RGB
        img_arr_SR_RGB = cv2.cvtColor(np.asarray(img_arr_SR_YUV,dtype=np.uint8),cv2.COLOR_YUV2RGB)
        img_arr_blur_RGB = cv2.cvtColor(img_arr_blur_YUV,cv2.COLOR_YUV2RGB)
        plt.subplot(1,3,1)
        io.imshow(img_arr_RGB)
        plt.subplot(1,3,2)
        io.imshow(img_arr_blur_RGB)
        plt.subplot(1,3,3)
        io.imshow(img_arr_SR_RGB)
        #PSNR
        print 'PSNR-originalVSbi: {0}'.format(measure.compare_psnr(im_true=img_arr_RGB,im_test=img_arr_blur_RGB,dynamic_range=255))
        print 'PSNR-originalVSsr: {0}'.format(measure.compare_psnr(im_true=img_arr_YUV,im_test=img_arr_SR_RGB,dynamic_range=255))
        print ""*40
        io.show()


if __name__ == "__main__":
    inference()